from typing import Optional, Callable, List
import os.path as osp
import os

import torch
from torch_geometric.data import Data
from torch_geometric.data import InMemoryDataset
from torch_geometric.data.storage import GlobalStorage
from torch_geometric.data.data import DataEdgeAttr, DataTensorAttr

from utils.data_utils import read_npz_file
from .load_data import circuit_parse_pyg
from .ordered_data import OrderedData

# 添加OrderedData, GlobalStorage, DataEdgeAttr和DataTensorAttr到安全全局变量列表
torch.serialization.add_safe_globals([OrderedData, GlobalStorage, DataEdgeAttr, DataTensorAttr])

class CircuitDataset(InMemoryDataset):
    r"""
    A variety of circuit graph datasets, *e.g.*, open-sourced benchmarks,
    random circuits.

    Args:
        root (string): Root directory where the dataset should be saved.
        args (object): The arguments specified by the main program.
        transform (callable, optional): A function/transform that takes in an
            :obj:`torch_geometric.data.Data` object and returns a transformed
            version. The data object will be transformed before every access.
            (default: :obj:`None`)
        pre_transform (callable, optional): A function/transform that takes in
            an :obj:`torch_geometric.data.Data` object and returns a
            transformed version. The data object will be transformed before
            being saved to disk. (default: :obj:`None`)
        pre_filter (callable, optional): A function that takes in an
            :obj:`torch_geometric.data.Data` object and returns a boolean
            value, indicating whether the data object should be included in the
            final dataset. (default: :obj:`None`)
    """

    def __init__(self, root, args, transform=None, pre_transform=None, pre_filter=None, circuit_file=None, label_file=None):
        """加载电路图数据集

        Args:
            root: 数据集根目录
            args: 参数对象
            transform: 转换函数
            pre_transform: 转换函数
            pre_filter: 过滤器
            circuit_file: 电路图数据文件名
            label_file: 标签数据文件名
        """
        # 首先确保args属性有所有必要的字段，避免processed_dir出错
        self.args = args
        
        # 为args添加默认属性，如果它们不存在的话
        default_attrs = {
            'use_edge_attr': 0,
            'reconv_skip_connection': 0,
            'predict_diff': 0,
            'diff_multiplier': 0,
            'no_node_cop': 0,
            'node_reconv': 0,
            'num_gate_types': 4,
            'dim_edge_feature': 4,
            'small_train': 0,
            'un_directed': 0
        }
        
        for attr, default_value in default_attrs.items():
            if not hasattr(self.args, attr):
                setattr(self.args, attr, default_value)
        
        # 设置电路文件和标签文件名
        self.circuit_file = circuit_file
        if self.circuit_file is None and hasattr(args, 'circuit_file') and args.circuit_file:
            self.circuit_file = args.circuit_file
        elif self.circuit_file is None:
            self.circuit_file = "benchmarks_circuits_graphs.npz"
            
        self.label_file = label_file
        if self.label_file is None and hasattr(args, 'label_file') and args.label_file:
            self.label_file = args.label_file
        elif self.label_file is None:
            self.label_file = "benchmarks_circuits_labels.npz"
            
        # 先调用父类初始化，再使用processed_dir
        super().__init__(root, transform, pre_transform, pre_filter)
        
        # 加载处理好的数据
        try:
            self.processed_path = osp.join(self.processed_dir, 'data.pt')
            
            if osp.exists(self.processed_path):
                self.data, self.slices = torch.load(self.processed_path)
                print(f"[INFO] 成功从{self.processed_path}加载处理好的数据")
            else:
                print(f"[WARNING] 找不到处理好的数据文件：{self.processed_path}")
                self.process()
        except Exception as e:
            print(f"[ERROR] 加载数据集时出错: {str(e)}")
            raise

    @property
    def raw_dir(self):
        return self.root

    @property
    def processed_dir(self):
        name = "{}_{}_{}_{}_{}_{}_{}_{}_{}_{}".format(int(self.args.use_edge_attr), int(self.args.reconv_skip_connection), int(self.args.predict_diff), self.args.diff_multiplier, int(self.args.no_node_cop), int(self.args.node_reconv), self.args.num_gate_types, self.args.dim_edge_feature, int(self.args.small_train), int(self.args.un_directed))
        return osp.join(self.root, name)

    @property
    def raw_file_names(self) -> List[str]:
        return [self.circuit_file, self.label_file]

    @property
    def raw_paths(self) -> List[str]:
        # 手动实现raw_paths，避免依赖于Dataset基类的实现
        # 返回电路文件和标签文件的完整路径
        circuit_path = osp.join(self.root, self.circuit_file)
        label_path = osp.join(self.root, self.label_file)
        return [circuit_path, label_path]

    @property
    def processed_file_names(self) -> str:
        return ['data.pt']

    def download(self):
        pass

    def _process_data(self):
        """直接处理数据而不依赖Dataset基类的处理流程"""
        try:
            # 获取单个电路名称（如果指定）
            single_circuit = self.args.circuit if hasattr(self.args, 'circuit') else ''
            
            # 尝试检测电路文件名
            circuit_path = self.raw_paths[0]
            label_path = self.raw_paths[1] if len(self.raw_paths) > 1 else None
            
            # 检查并设置缺失的属性默认值
            if not hasattr(self.args, 'logic_diff_embedding'):
                self.args.logic_diff_embedding = 'positional'
                print("[INFO] 设置默认的logic_diff_embedding='positional'")
                
            # 如果找不到默认电路文件，尝试在数据目录中搜索
            if not os.path.exists(circuit_path):
                print(f"[WARNING] 找不到默认电路文件: {circuit_path}")
                print(f"[INFO] 正在搜索目录 {self.root} 中的其他电路文件...")
                
                circuit_files = []
                for file in os.listdir(self.root):
                    if file.endswith('.npz') and ('circuit' in file.lower() or 'benchmark' in file.lower()):
                        self.circuit_file = file
                        circuit_path = os.path.join(self.root, file)
                        circuit_files.append(file)
                        print(f"[INFO] 找到电路文件: {circuit_path}")
                        break
                
                if not circuit_files:
                    print(f"[WARNING] 在 {self.root} 中找不到电路文件")
                    return []
            
            if os.path.exists(circuit_path):
                try:
                    print(f"[INFO] 正在加载电路文件: {circuit_path}")
                    circuits = read_npz_file(self.circuit_file, self.root)['circuits'].item()
                    
                    print(f"[INFO] 尝试加载标签文件: {label_path}")
                    if self.args.no_labels:
                        print("[INFO] 使用空标签 (no_labels=True)")
                        labels = {}
                        for cir_idx, cir_name in enumerate(circuits):
                            label = [0] * len(circuits[cir_name]["x"])
                            labels[cir_name] = {"y": label}
                    elif label_path and os.path.exists(label_path):
                        labels = read_npz_file(self.label_file, self.root)['labels'].item()
                    else:
                        print("[WARNING] 标签文件不存在，使用空标签")
                        labels = {}
                        for cir_idx, cir_name in enumerate(circuits):
                            label = [0] * len(circuits[cir_name]["x"])
                            labels[cir_name] = {"y": label}
                    
                    if self.args.small_train:
                        subset = 10
                    
                    data_list = []
                    print(f"[INFO] 可用电路: {list(circuits.keys())}")
                    
                    # 如果指定了单个电路，只处理该电路
                    if single_circuit:
                        print(f"[INFO] 指定测试单个电路: {single_circuit}")
                        if single_circuit in circuits:
                            cir_name = single_circuit
                            print(f'[INFO] 处理电路: {cir_name}')
                            x = circuits[cir_name]["x"]
                            edge_index = circuits[cir_name]["edge_index"]
                            
                            y = labels[cir_name]["y"] if cir_name in labels else [0] * len(x)
                            
                            graph = circuit_parse_pyg(x, edge_index, y, self.args.use_edge_attr, \
                                self.args.reconv_skip_connection, self.args.logic_diff_embedding, self.args.predict_diff, \
                                self.args.diff_multiplier, self.args.no_node_cop, self.args.node_reconv, self.args.un_directed, self.args.num_gate_types, self.args.dim_edge_feature)
                            graph.len = len(x)
                            graph.name = cir_name
                            data_list.append(graph)
                            print(f"[INFO] 成功处理电路: {cir_name}")
                        else:
                            # 尝试查找相似名称的电路
                            found_match = False
                            closest_match = None
                            for cir_name in circuits.keys():
                                # 检查电路名是否为指定电路的一部分，或者指定电路是电路名的一部分
                                if single_circuit.lower() in cir_name.lower() or cir_name.lower() in single_circuit.lower():
                                    print(f"[INFO] 找到近似匹配电路: {cir_name} (用户指定: {single_circuit})")
                                    closest_match = cir_name
                                    found_match = True
                                    break
                                    
                            if found_match:
                                # 处理找到的近似匹配电路
                                cir_name = closest_match
                                print(f'[INFO] 处理近似匹配电路: {cir_name} (用户指定: {single_circuit})')
                                x = circuits[cir_name]["x"]
                                edge_index = circuits[cir_name]["edge_index"]
                                
                                y = labels[cir_name]["y"] if cir_name in labels else [0] * len(x)
                                
                                graph = circuit_parse_pyg(x, edge_index, y, self.args.use_edge_attr, \
                                    self.args.reconv_skip_connection, self.args.logic_diff_embedding, self.args.predict_diff, \
                                    self.args.diff_multiplier, self.args.no_node_cop, self.args.node_reconv, self.args.un_directed, self.args.num_gate_types, self.args.dim_edge_feature)
                                graph.len = len(x)
                                # 保存用户指定的原始电路名，以保持一致性
                                graph.name = single_circuit
                                graph.original_name = cir_name
                                data_list.append(graph)
                                print(f"[INFO] 成功处理近似匹配电路: {cir_name} (保存为: {single_circuit})")
                            else:
                                print(f"[ERROR] 找不到指定的电路: {single_circuit}")
                                print(f"[INFO] 可用电路列表: {list(circuits.keys())}")
                    else:
                        # 处理所有电路
                        for cir_idx, cir_name in enumerate(circuits):
                            print(f'[INFO] 处理电路: {cir_name}')
                            x = circuits[cir_name]["x"]
                            edge_index = circuits[cir_name]["edge_index"]
                            
                            y = labels[cir_name]["y"] if cir_name in labels else [0] * len(x)
                            
                            graph = circuit_parse_pyg(x, edge_index, y, self.args.use_edge_attr, \
                                self.args.reconv_skip_connection, self.args.logic_diff_embedding, self.args.predict_diff, \
                                self.args.diff_multiplier, self.args.no_node_cop, self.args.node_reconv, self.args.un_directed, self.args.num_gate_types, self.args.dim_edge_feature)
                            graph.len = len(x)
                            graph.name = cir_name
                            data_list.append(graph)
                            if self.args.small_train and cir_idx > subset:
                                break
                    
                    print(f"[INFO] 成功处理了 {len(data_list)} 个电路")
                    return data_list
                except Exception as e:
                    print(f"[ERROR] 处理电路文件时出错: {str(e)}")
                    import traceback
                    traceback.print_exc()
                    return []
            else:
                print(f"[ERROR] 电路文件不存在: {circuit_path}")
                return []
        except Exception as e:
            print(f"[ERROR] 处理数据时出错: {str(e)}")
            import traceback
            traceback.print_exc()
            return []

    def process(self):
        data_list = self._process_data()
        if data_list:
            data, slices = self.collate(data_list)
            torch.save((data, slices), self.processed_paths[0])
        else:
            # 创建一个伪数据文件，避免重复处理
            print("[WARNING] 创建空数据集，避免重复处理")
            os.makedirs(os.path.dirname(self.processed_paths[0]), exist_ok=True)
            
            # 确保dim_node_feature是整数，防止出现类型错误
            try:
                dim_node_feature = int(self.args.dim_node_feature) if hasattr(self.args, 'dim_node_feature') else 4
            except (ValueError, TypeError):
                print("[WARNING] dim_node_feature转换为整数失败，使用默认值4")
                dim_node_feature = 4
            
            # 创建一个最小的数据集，确保使用元组格式
            print(f"[DEBUG] 创建空数据集，node_feature维度: {dim_node_feature}")
            x = torch.zeros((1, dim_node_feature))
            edge_index = torch.zeros((2, 0), dtype=torch.long)
            y = torch.zeros(1)
            
            from .ordered_data import OrderedData
            data = OrderedData(x=x, edge_index=edge_index, y=y)
            data.len = 1
            data.name = "empty"
            slices = {'x': torch.tensor([0, 1]), 'edge_index': torch.tensor([0, 0]), 'y': torch.tensor([0, 1])}
            
            torch.save((data, slices), self.processed_paths[0])

    def __repr__(self) -> str:
        return f'{self.name}({len(self)})'